{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jacobians in navlie\n",
    "\n",
    "As you may know, many state estimation algorithm require access to process model and measurement model Jacobians, with respect to the state and sometimes other inputs. For states belonging to Lie groups, algorithms will require _Lie Jacobians_, which differ from traditional derivatives as they conform to the constraints of the group. The abstraction provided by the $\\oplus$ and $\\ominus$ operators (implemented with `State.plus` and `State.minus` respectively) allow for a generic definition of a derivative:\n",
    "\n",
    "$$\n",
    "\\left.\\frac{D f(\\mathcal{X})}{D \\mathcal{X}}\\right|_{\\bar{\\mathcal{X}} }\\triangleq \\left.\\frac{\\partial f(\\bar{\\mathcal{X}} \\oplus \\delta \\mathbf{x}) \\ominus f(\\bar{\\mathcal{X}})}{\\partial \\delta \\mathbf{x}}\\right|_{\\delta \\mathbf{x} = \\mathbf{0}},\n",
    "$$\n",
    "\n",
    "which can be shown to fall back to a traditional derivatives when  $\\oplus$ and $\\ominus$ are defined to be regular addition/subtraction. This derivative definition is used universally throughout navlie, and roughly follows what is done in the [Micro Lie Theory paper](https://arxiv.org/pdf/1812.01537.pdf). In that reference, seperate definitions are given for \"left\" and \"right\" derivatives, whereas we have aggregated them into a single definition, with left and right derivatives naturally arising from the choice of $\\oplus$ and $\\ominus$ operators.\n",
    "\n",
    "If you dont want to worry about this, the good news is that navlie computes Lie Jacobians for you by default using finite difference. However, finite difference can have some drawbacks, such as being computationally expensive and less accurate than analytic derivatives. In this notebook, we will show you how to use analytic derivatives in navlie, which offer the maximum accuracy and speed.\n",
    "\n",
    "## Jacobians - Traditional Approach\n",
    "Recall the traditional approach to the previous example. We had defined the state to be $\\mathbf{x} = [\\theta, x, y]^T$ and the process model to be\n",
    "\n",
    "$$ \n",
    "\\begin{aligned}\n",
    "\\theta_{k+1} &= \\theta_k + \\omega_k \\Delta t \\\\\n",
    "x_{k+1} &= x_k + v_k \\cos(\\theta_k) \\Delta t \\\\\n",
    "y_{k+1} &= y_k + v_k \\sin(\\theta_k) \\Delta t\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $\\omega_k$ is the angular velocity and $v_k$ is the linear velocity. Since the state is just a regular vector, Lie Jacobians fall back to regular Jacobians, and standard derivative techniques lead to the following expressions for the process model Jacobian with respect to the state\n",
    "\n",
    "$$\n",
    "\\mathbf{F} := \\frac{\\partial f(\\mathbf{x}_k, \\mathbf{u}_k)}{\\partial \\mathbf{x}_k} = \\begin{bmatrix} 1 & 0 & 0 \\\\ -v_k \\sin(\\theta_k) \\Delta t & 1 & 0 \\\\ v_k \\cos(\\theta_k) \\Delta t & 0 & 1 \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "To implement this Jacobian in navlie, all we need to do is override the `jacobian()` method in our process model, and it will get used automatically by the estimation algorithms. Adding this to our process model from before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import navlie as nav\n",
    "import numpy as np\n",
    "\n",
    "from navlie.lib import VectorInput, VectorState\n",
    "\n",
    "Q = np.eye(2) * 0.1**2 # Input noise covariance with 0.1 m/s of standard deviation\n",
    "\n",
    "class WheeledRobot(nav.ProcessModel):\n",
    "    def __init__(self, input_covariance):\n",
    "        self.Q = input_covariance\n",
    "\n",
    "    def evaluate(self, x: VectorState, u: VectorInput, dt: float) -> VectorState:\n",
    "        x_next = x.copy()\n",
    "        x_next.value[0] += u.value[0] * dt\n",
    "        x_next.value[1] += u.value[1] * dt * np.cos(x.value[0])\n",
    "        x_next.value[2] += u.value[1] * dt * np.sin(x.value[0])\n",
    "        return x_next\n",
    "\n",
    "    def input_covariance(self, x: VectorState, u: VectorInput, dt: float) -> np.ndarray:\n",
    "        return self.Q\n",
    "    \n",
    "    def jacobian(self, x:VectorState, u: VectorInput, dt: float) -> np.ndarray:\n",
    "        F = np.eye(3)\n",
    "        F[1, 0] = -u.value[1] * dt * np.sin(x.value[0])\n",
    "        F[2, 0] = u.value[1] * dt * np.cos(x.value[0])\n",
    "        return F\n",
    "\n",
    "process_model = WheeledRobot(Q) # instantiate it"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets just double check that we did everything correctly by comparing with finite difference. All process models inherit the `jacobian_fd()` method, which computes the Jacobian using finite difference. We can use this to compare with our analytic Jacobian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Analyical:\n",
      "[[ 1.          0.          0.        ]\n",
      " [-0.01682942  1.          0.        ]\n",
      " [ 0.01080605  0.          1.        ]]\n",
      "\n",
      "Finite difference:\n",
      "[[ 1.          0.          0.        ]\n",
      " [-0.01682943  1.          0.        ]\n",
      " [ 0.01080604  0.          1.        ]]\n"
     ]
    }
   ],
   "source": [
    "x = VectorState([1,2,3]) \n",
    "u = VectorInput([0.1, 0.2]) \n",
    "dt = 0.1 \n",
    "\n",
    "print(\"Analyical:\")\n",
    "print(process_model.jacobian(x, u, dt))\n",
    "\n",
    "print(\"\\nFinite difference:\")\n",
    "print(process_model.jacobian_fd(x, u, dt))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Jacobians match almost perfectly, but differ slightly due to errors in finite difference. This is expected, as finite difference is only an approximation. Nevertheless, finite difference is generally sufficiently accurate for most applications.\n",
    "\n",
    "<div class=\"alert alert-block alert-warning\">\n",
    "Note: The `jacobian` methods must *always* return a 2D array.\n",
    "</div>\n",
    "\n",
    "\n",
    "\n",
    "Moving on to the measurement model, which was previously defined to be \n",
    "\n",
    "$$\n",
    "\\mathbf{y}_k = || \\mathbf{r}_k - \\boldsymbol{\\ell}^{(i)} ||  \n",
    "$$ \n",
    "\n",
    "where $\\mathbf{r}_k$ is the robot's position and $\\boldsymbol{\\ell}^{(i)}$ is the $i$th landmark. The measurement model Jacobian with respect to the state is\n",
    "\n",
    "$$ \n",
    "\\mathbf{G} := \\frac{\\partial g(\\mathbf{x}_k)}{\\partial \\mathbf{x}_k} = \\begin{bmatrix} 0 & \\frac{(\\mathbf{r}_k - \\boldsymbol{\\ell}^{(i)})^T}{||(\\mathbf{r}_k - \\boldsymbol{\\ell}^{(i)})||} \\end{bmatrix},\n",
    "$$\n",
    "\n",
    "and we can implement this in navlie by again overriding the `jacobian()` method in our measurement model. Adding this to our measurement model from before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Analyical:\n",
      "[[0.         0.70710678 0.70710678]]\n",
      "\n",
      "Finite difference:\n",
      "[[0.         0.70710696 0.70710696]]\n"
     ]
    }
   ],
   "source": [
    "class RangeToLandmark(nav.MeasurementModel):\n",
    "    def __init__(self, landmark_position: np.ndarray):\n",
    "        self.landmark_position = landmark_position\n",
    "\n",
    "    def evaluate(self, x: VectorState) -> np.ndarray:\n",
    "        return np.linalg.norm(x.value[1:] - self.landmark_position)\n",
    "    \n",
    "    def covariance(self, x: VectorState) -> np.ndarray:\n",
    "        return 0.1**2\n",
    "    \n",
    "    def jacobian(self, x: VectorState) -> np.ndarray:\n",
    "        G = np.zeros((1, 3))\n",
    "        r = x.value[1:]\n",
    "        G[0,1:] = (r - self.landmark_position) / np.linalg.norm(r - self.landmark_position)\n",
    "        return G\n",
    "    \n",
    "meas_model = RangeToLandmark(np.array([1, 2]))\n",
    "\n",
    "print(\"\\nAnalyical:\")\n",
    "print(meas_model.jacobian(x))\n",
    "\n",
    "print(\"\\nFinite difference:\")\n",
    "print(meas_model.jacobian_fd(x))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can again see that the results match nicely."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Jacobians - Lie Group Approach \n",
    "\n",
    "Now, lets see how to implement analytical Jacobians when states belong to Lie groups. In the previous example the state was $\\mathbf{T} \\in SE(2)$ and the process model was\n",
    "\n",
    "$$\n",
    "\\mathbf{T}_{k+1} = \\mathbf{T}_k \\exp(\\boldsymbol{\\varpi}^\\wedge_k \\Delta t).\n",
    "$$\n",
    "\n",
    "To derive the Jacobian, we can \"perturb\" both sides of the equation and manipulate. This is a common technique for deriving Lie Jacobians, and for computing matrix-vector derivatives in general. For more details, we recommend reading [State Estimation for Robotics by Tim Barfoot](http://asrl.utias.utoronto.ca/~tdb/bib/barfoot_ser17.pdf).\n",
    "\n",
    "$$ \n",
    "\\begin{aligned}\n",
    "\\mathbf{T}_{k+1} &= \\mathbf{T}_k \\exp(\\boldsymbol{\\varpi}^\\wedge_k \\Delta t) \\\\\n",
    "\\bar{\\mathbf{T}}_{k+1} \\exp(\\delta \\boldsymbol{\\xi}_{k+1}^\\wedge) &= \\bar{\\mathbf{T}}_{k} \\exp(\\delta \\boldsymbol{\\xi}_{k}^\\wedge) \\exp(\\boldsymbol{\\varpi}^\\wedge_k \\Delta t) \\\\\n",
    " \\exp(\\delta \\boldsymbol{\\xi}_{k+1}^\\wedge) &= \\underbrace{\\bar{\\mathbf{T}}_{k+1}^{-1} \\bar{\\mathbf{T}}_{k}}_{\\exp(\\boldsymbol{\\varpi}^\\wedge \\Delta t)^{-1}} \\exp(\\delta \\boldsymbol{\\xi}_{k}^\\wedge) \\exp(\\boldsymbol{\\varpi}^\\wedge_k \\Delta t) \\\\\n",
    " \\exp(\\delta \\boldsymbol{\\xi}_{k+1}^\\wedge) &= \\exp((\\mathbf{Ad}(\\exp(\\boldsymbol{\\varpi}^\\wedge \\Delta t)^{-1}) \\delta \\boldsymbol{\\xi}_{k})^\\wedge) \\\\\n",
    " \\delta \\boldsymbol{\\xi}_{k+1} &= \\underbrace{\\mathbf{Ad}(\\exp(\\boldsymbol{\\varpi}^\\wedge \\Delta t)^{-1})}_{\\mathbf{F}} \\delta \\boldsymbol{\\xi}_{k}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "There, we used the _adjoint matrix_ $\\mathbf{Ad}(\\cdot)$, to invoke the identity $\\mathbf{X}^{-1} \\exp(\\boldsymbol{\\xi}^\\wedge) \\mathbf{X} = \\exp(\\mathbf{Ad}(\\mathbf{X}) \\boldsymbol{\\xi}^\\wedge)$, which is true for any $\\mathbf{X} \\in SE(2)$. The adjoint matrix for $SE(2)$ is given by \n",
    "\n",
    "$$ \n",
    "\\mathbf{Ad}(\\mathbf{T}) = \\begin{bmatrix} 1 & \\mathbf{0} \\\\ - \\boldsymbol{\\Omega} \\mathbf{r} & \\mathbf{C} \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "where $\\boldsymbol{\\Omega} = \\begin{bmatrix} 0 & -1 \\\\  1 \n",
    "&0 \\end{bmatrix}$.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "Note: in this Jacobian derivation, we perturbed the state $\\mathbf{T} = \\bar{\\mathbf{T}} \\exp(\\delta \\boldsymbol{\\xi}^\\wedge)$ \"on the right\" because that corresponds to what was implemented in the `plus()` method of our `SE2State` class. It is important to be consistent here for everything to work.\n",
    "</div>\n",
    "\n",
    "Now, we can implement this Jacobian in navlie by overriding the `jacobian()` method in our process model. Adding this to our process model from before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Analyical:\n",
      "[[ 1.          0.          0.        ]\n",
      " [ 0.0394695   0.92106099  0.38941834]\n",
      " [ 0.19470917 -0.38941834  0.92106099]]\n",
      "\n",
      "Finite difference:\n",
      "[[ 1.00000000e+00 -1.10747782e-11 -1.10747782e-11]\n",
      " [ 3.94695038e-02  9.21060995e-01  3.89418343e-01]\n",
      " [ 1.94709171e-01 -3.89418342e-01  9.21060994e-01]]\n"
     ]
    }
   ],
   "source": [
    "from navlie.lib import SE2State, VectorInput\n",
    "from scipy.linalg import expm\n",
    "\n",
    "\n",
    "def wedge_se2(x:np):\n",
    "    return np.array([[   0, -x[0], x[1]],\n",
    "                     [x[0],     0, x[2]], \n",
    "                     [   0,     0,    0]])\n",
    "\n",
    "def adjoint_se2(T:np.ndarray):\n",
    "    C = T[:2, :2]\n",
    "    r = T[:2, 2]\n",
    "    Omega = np.array([[0, -1], [1, 0]])\n",
    "    Ad = np.zeros((3,3))\n",
    "    Ad[0,0] = 1\n",
    "    Ad[1:,1:] = C\n",
    "    Ad[1:,0] = - Omega @ r\n",
    "    return Ad\n",
    "\n",
    "\n",
    "class WheeledRobotSE2(nav.ProcessModel):\n",
    "    def __init__(self, input_covariance_matrix):\n",
    "        self.Q = input_covariance_matrix\n",
    "\n",
    "    def evaluate(self, x:SE2State, u:VectorInput, dt:float):\n",
    "        u = np.array([u.value[0], u.value[1], 0])\n",
    "        x_next = x.copy()\n",
    "        x_next.value = x.value @ expm(wedge_se2(u * dt))\n",
    "        return x_next\n",
    "    def input_covariance(self, x:SE2State, u:VectorInput, dt:float):\n",
    "        return self.Q\n",
    "\n",
    "    def jacobian(self, x:SE2State, u:VectorInput, dt:float):\n",
    "        u = np.array([u.value[0], u.value[1], 0])\n",
    "        return adjoint_se2(expm(-wedge_se2(u * dt)))\n",
    "\n",
    "Q = np.eye(2) * 0.1**2\n",
    "process_model = WheeledRobotSE2(Q)\n",
    "x = SE2State(expm(wedge_se2(np.array([1,2,3]))))\n",
    "u = VectorInput([4, 2])\n",
    "dt = 0.1\n",
    "\n",
    "print(\"\\nAnalyical:\")\n",
    "print(process_model.jacobian(x, u, dt))\n",
    "\n",
    "print(\"\\nFinite difference:\")\n",
    "print(process_model.jacobian_fd(x, u, dt))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that when using Lie groups, our Jacobian no longer has dependence on the state itself. This can be a tangible advantage when the state estimate has high uncertainty, where using a traditional approach can result in excessive linearization errors when the state estimate is far from the true value."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
